#include "m8515def.inc"
#include "Commands.inc"

#define BOOTSTART	THIRDBOOTSTART

#define ATMEL		0x1e

#define	PAGESIZEB	(PAGESIZE << 1)

/*********************************************************/
/*                       PORTS							 */
/*********************************************************/

#define GET_PINS(port)	(port-2)
#define	GET_DDR(port)	(port-1)

#define DATA_PORT	PORTD
#define	DATA_DDR	GET_DDR(DATA_PORT)
#define	DATA_PINS	GET_PINS(DATA_PORT)

#define CTRL_PORT	PORTB
#define	CTRL_DDR	GET_DDR(CTRL_PORT)
#define	CTRL_PINS	GET_PINS(CTRL_PORT)

#define	RD_PIN		4
#define	RXF_PIN		3
#define WR_PIN		2
#define	TXE_PIN		1
#define	LDR_EN_PIN	0

/*********************************************************/
/*						REGISTERS						 */
/*********************************************************/

#define	prodl	r0
#define	prodh	r1	

#define var1	r16
#define	var2	r17
#define	var3	r18
#define var4	r19

#define	spmcrval	var3

#define	ctr1	XL
#define	ctr2	XH

#define	locvar1	r20
#define	locvar2	r21

#define	arg1	r24
#define	arg2	r25

#define	result	arg1

/*********************************************************/
/*					   COMMUNICATION				 	 */
/*********************************************************/

.MACRO	Read
		rcall	ReadUSB
.ENDM

.MACRO	Write
		rcall	WriteUSB
.ENDM

/*********************************************************/
/*					    BOOTLOADER						 */
/*********************************************************/

.ORG	0

start:	
	//	rjmp	boot

.ORG	BOOTSTART	

		rjmp	boot

#include "FusionBootInfo.asm"

boot:	// Bootloader
		
		// Disable Interrupts
		cli
		
		// Check, whether to enter bootloader.
		cbi		CTRL_DDR,	LDR_EN_PIN
		sbi		CTRL_PORT,	LDR_EN_PIN
		nop
		sbis	CTRL_PINS,	LDR_EN_PIN
		rjmp	enter
		cbi		CTRL_PORT,	LDR_EN_PIN
		rjmp	start

enter:
		// Init Stack
		ldi		var1,	low(RAMEND)
		out		SPL,	var1
		ldi		var1,	high(RAMEND)
		out		SPH,	var1

		// Init	DATA_PORT
		clr		var1
		out		DATA_DDR,	var1
		ser		var1
		out		DATA_PORT,	var1
		
		// Init CTRL_PORT
		ldi		var1,	(1<<RD_PIN)|(1<<WR_PIN)|(1<<RXF_PIN)|(1<<TXE_PIN)
		out		CTRL_PORT,	var1
		ldi		var1,	(1<<RD_PIN)|(1<<WR_PIN)
		out		CTRL_DDR,	var1

waitEn:	// Wait for Enable command.
		Read
		cpi		result,	cmdEnable
		breq	Enable
		ldi		arg1,	rNotEnabled
		Write
		rjmp	waitEn
		
		// Send rOK	and enter main loop.
Enable:	ldi		arg1,	rOK
		Write

loop:	Read
		cpi		result,	cmdEnable
		breq	Enable
		
		cpi		result,	cmdWriteProgramMemoryPage
		breq	_WriteProgramMemoryPage
		
		cpi		result,	cmdReadProgramMemoryPage
		breq	ReadProgramMemoryPage
		
		cpi		result, cmdEraseProgramMemoryPage
		breq	_EraseProgramMemoryPage

		cpi		result, cmdGetFusionInfo
		breq	GetFusionInfo
		
		cpi		result,	cmdGetDeviceInfo		
		breq 	GetDeviceInfo

		ldi		arg1,	rNotSupported
		Write
		rjmp	loop

/*********************************************************/
/*					 BRANCH	HELPERS						 */
/*********************************************************/
_WriteProgramMemoryPage:
		rjmp	WriteProgramMemoryPage

_EraseProgramMemoryPage:
		rjmp	EraseProgramMemoryPage

/*********************************************************/
/*					  		INFO						 */
/*********************************************************/

/* Returns fusion info */
GetFusionInfo:
		push	ZL
		push	ZH
		ldi		arg1,	rOK
		Write

		ldi		ctr1,	FUSION_BOOT_INFO_SIZE // Counter
		mov		arg1,	ctr1
		Write
		
		ldi		ZL,	low(FUSION_BOOT_INFO)
		ldi		ZH, high(FUSION_BOOT_INFO)
		clr		var1	// Control Sum
  
  gfi1:	lpm		arg1,	Z+
  		eor		var1,	arg1
		Write
		subi	ctr1,	1
		brne	gfi1

		mov		arg1,	var1
		Write

		pop		ZH
		pop		ZL
		rjmp	loop

/* Returns device info */
GetDeviceInfo:  
		ldi		arg1,	rOK
		Write
		ldi		arg1,	4	// Data Length
		Write
		ldi		arg1,	ATMEL
		Write
		ldi		arg1,	SIGNATURE_000
		Write
		ldi		arg1,	SIGNATURE_001
		Write
		ldi		arg1,	SIGNATURE_002
		Write
		ldi		arg1,	(ATMEL ^ SIGNATURE_000 ^ SIGNATURE_001 ^ SIGNATURE_002)
		Write
		rjmp	loop

/*********************************************************/
/*					   PROGRAM MEMORY					 */
/*********************************************************/

/* Reads a page of program memory. */
ReadProgramMemoryPage:
		push	ZL
		push	ZH

		// Write OK
		ldi		arg1,	rOK
		Write

		// Read Page Number
		Read
		mov		var3,	result	// Page address - low byte
		Read
		mov		var4,	result	// Page address - high byte
		Read	
		mov		var2,	result	// Control sum
		
		// Calculate control sum
		ldi		var1,	cmdReadProgramMemoryPage
		eor		var1,	var3
		eor		var1,	var4
		cp		var1,	var2
		breq	rpmp1

		ldi		arg1,	rControlSumError
		Write
		rjmp	rpmpx
  
  rpmp1:
		ldi		arg1,	rOK
		Write
		
		clr		var1	// Control sum

		// Calculate address
		ldi		ctr1,	low(PAGESIZEB)
		ldi		ctr2,	high(PAGESIZEB)
		mul		ctr1,	var3
		movw	ZL,		prodl
		mul		ctr1,	var4
		add		ZL,		prodl
		adc		ZH,		prodh
		mul		ctr2,	arg1
		add		ZL,		prodl
		adc		ZH,		prodh
  rpmp2:
		lpm		arg1,	Z+
		eor		var1,	arg1
		Write
		sbiw	ctr1,	1
		brne	rpmp2
		
		mov		arg1,	var1
		Write		
  rpmpx:
		pop		ZH
		pop		ZL
		rjmp	loop


/* Writes program memory page. */
WriteProgramMemoryPage:
		push	ZL
		push	ZH
		
		// Write OK
		ldi		arg1,	rOK
		Write

		// Read Page Number
		Read
		mov		var3,	result	// Page address - low byte
		Read
		mov		var4,	result	// Page address - high byte
		
		// Calculate control sum
		ldi		var1,	cmdWriteProgramMemoryPage
		eor		var1,	var3
		eor		var1,	var4

		// Load RAM pointer
		ldi		YL,		low(pageBuffer)
		ldi		YH,		high(pageBuffer)

		// Load Counter
		ldi		ctr1,	low(PAGESIZEB)
		ldi		ctr2,	high(PAGESIZEB)

  wpmp1: //Read Page
  		Read
		eor		var1,	result
		st		Y+,		result
		sbiw	ctr1,	1
		brne	wpmp1
		
		Read	result
		cp		var1,	result
		breq	wpmp2

		ldi		arg1,	rControlSumError
		Write
		rjmp	wpmpx

  wpmp2: // Check page number
  		ldi		var1,	high(BOOTSTART/PAGESIZE)	
		cp		var1,	var4
		brlo	wpmp3
		cpi		var3,	low(BOOTSTART/PAGESIZE)
		brcc	wpmp3
		rjmp	wpmp4	
		
  wpmp3: // Error: Non-writable page
  		ldi		arg1,	rNonWritable
		Write
		rjmp	wpmpx		

  wpmp4: // Calculate address
		ldi		ctr1,	low(PAGESIZEB)
		ldi		ctr2,	high(PAGESIZEB)
		mul		ctr1,	var3
		movw	ZL,		prodl
		mul		ctr1,	var4
		add		ZL,		prodl
		adc		ZH,		prodh
		mul		ctr2,	arg1
		add		ZL,		prodl
		adc		ZH,		prodh
		  	
		// Load RAM pointer
		ldi		YL,		low(pageBuffer)
		ldi		YH,		high(pageBuffer)
	
		// Perform Programming

		// page erase
		ldi 	spmcrval, (1<<PGERS) | (1<<SPMEN)
		rcall	Do_spm

		// re-enable the RWW section
		ldi 	spmcrval, (1<<RWWSRE) | (1<<SPMEN)
		rcall	Do_spm		

		// transfer data from RAM to Flash page buffer
  wpmp5:
		ld 		r0, Y+
		ld 		r1, Y+
		ldi 	spmcrval, (1<<SPMEN)
		rcall	Do_spm
		adiw 	ZH:ZL,	2
		sbiw 	ctr1,	2 
		brne 	wpmp5
		// execute page write
		subi 	ZL, 	low(PAGESIZEB) //restore pointer
		sbci 	ZH, 	high(PAGESIZEB)
		ldi 	spmcrval, (1<<PGWRT) | (1<<SPMEN)
		rcall	Do_spm
		// re-enable the RWW section
		ldi 	spmcrval, (1<<RWWSRE) | (1<<SPMEN)
		rcall	Do_spm

		// read back and check, optional
		ldi 	ctr1, 	low(PAGESIZEB) 	//init loop variable
		ldi 	ctr2, 	high(PAGESIZEB)
		subi 	YL, 	low(PAGESIZEB) 	//restore pointer
		sbci 	YH, 	high(PAGESIZEB)
  wpmp6:
		lpm 	r0, 	Z+
		ld 		r1, 	Y+
		cpse 	r0, 	r1
		rjmp 	wpmp7
		sbiw 	ctr1,	1
		brne 	wpmp6

		ldi		arg1,	rOK
		Write
		rjmp	wpmpx

  wpmp7: // Write error handler
  		ldi		arg1,	rWriteError
		Write

  wpmpx:
		pop		ZH
		pop		ZL
		rjmp	loop

/* Erases Program Memory Page */
EraseProgramMemoryPage:
		push	ZL
		push	ZH

		// Write OK
		ldi		arg1,	rOK
		Write

		// Read Page Number
		Read
		mov		var3,	result	// Page address - low byte
		Read
		mov		var4,	result	// Page address - high byte
		Read	
		mov		var2,	result	// Control sum
		
		// Calculate control sum
		ldi		var1,	cmdEraseProgramMemoryPage
		eor		var1,	var3
		eor		var1,	var4
		cp		var1,	var2
		breq	epmp1

		ldi		arg1,	rControlSumError
		Write
		rjmp	epmpx

epmp1:	// Check page number
  		ldi		var1,	high(BOOTSTART/PAGESIZE)	
		cp		var1,	var4
		brlo	epmp2
		cpi		var3,	low(BOOTSTART/PAGESIZE)
		brcc	epmp2
		rjmp	epmp3

epmp2:	// Error: Non-writable page
  		ldi		arg1,	rNonWritable
		Write
		rjmp	epmpx
		
epmp3:	// Calculate address
		ldi		ctr1,	low(PAGESIZEB)
		ldi		ctr2,	high(PAGESIZEB)
		mul		ctr1,	var3
		movw	ZL,		prodl
		mul		ctr1,	var4
		add		ZL,		prodl
		adc		ZH,		prodh
		mul		ctr2,	arg1
		add		ZL,		prodl
		adc		ZH,		prodh

		// page erase
		ldi 	spmcrval, (1<<PGERS) | (1<<SPMEN)
		rcall	Do_spm

		// re-enable the RWW section
		ldi 	spmcrval, (1<<RWWSRE) | (1<<SPMEN)
		rcall	Do_spm

		ldi		arg1,	rOK
		Write
epmpx:
		pop		ZH
		pop		ZL
		rjmp	loop

Do_spm:
	// check for previous SPM complete
  	Wait_spm:
		in var1, SPMCR
		sbrc var1, SPMEN
		rjmp Wait_spm
	// check that no EEPROM write access is present
	Wait_ee:
		sbic EECR, EEWE
		rjmp Wait_ee
		// SPM timed sequence
		out SPMCR, spmcrval
		spm
		ret

/*********************************************************/
/*					  USB COMMUNICATION					 */
/*********************************************************/

/* Reads a byte from USB. */
ReadUSB: 
		sbic	CTRL_PINS,	RXF_PIN
		rjmp	ReadUSB

		clr		locvar1
		out		DATA_DDR,	locvar1

		cbi		CTRL_PORT,	RD_PIN
		nop		// This operation is required, because the MCU need one cycle to read from the port.
		in		result,	DATA_PINS
		sbi		CTRL_PORT,	RD_PIN
		ret

/* Writes a byte to USB. */
WriteUSB:
		sbic	CTRL_PINS,	TXE_PIN
		rjmp	WriteUSB

		ser		locvar1
		out		DATA_DDR,	locvar1

		out		DATA_PORT,	arg1
		cbi		CTRL_PORT,	WR_PIN
		sbi		CTRL_PORT,	WR_PIN
		ret

/*********************************************************/
/*							RAM							 */
/*********************************************************/

.dseg

pageBuffer:
		.byte	PAGESIZEB

.cseg
